import os

import numpy as np
from sklearn.model_selection import train_test_split

# def splitTrainTest(all_subjects):
#     '''
#     看github上，issues中作者的回答中，测试集是这几个
#     BTCV的CT的测试数据量为1/9/30/32/33/39，
#     CHAOS的MR数据为1/13/32/38/，其他数据量用于训练。
#     '''
#     test_list=['1','13','32','38']
#     train_list=[]
#     for subject in all_subjects:
#         if subject not in test_list:
#             train_list.append(subject)
#     return train_list,test_list
def splitTrainTest(all_subjects):
    '''
    看github上，issues中作者的回答中，测试集是这几个
    BTCV的CT的测试数据量为1/9/30/32/33/39，
    CHAOS的MR数据为1/13/32/38/，其他数据量用于训练。
    '''
    train_list,test_list=train_test_split(all_subjects,test_size=0.2,random_state=100)
    print(len(train_list),len(test_list))
    return train_list,test_list
def getTrainTest(source_path,target_path,txt_save_path='./'):
    '''
    将source和target的文件名保存到txt文件，用于后续dataset加载数据
    将target分为80%训练和20%测试
    '''
    #------------先处理源域的数据-------------#
    source_subjects_list=os.listdir(source_path)
    source_npz_name_list=[]
    for subject in source_subjects_list:
        for npz_name in os.listdir(os.path.join(source_path,subject)):
            abs_npz_name=os.path.join(source_path,subject,npz_name)
            source_npz_name_list.append(abs_npz_name)
    np.savetxt(os.path.join(txt_save_path, 'source.txt'), np.array(source_npz_name_list), fmt ="%s")
    #------------处理target数据集-------#
    all_target_subjects_list=os.listdir(target_path)
    target_train_subjects_list,target_test_subjects_list=splitTrainTest(all_target_subjects_list)
    #target的训练集
    target_train_npz_name_list = []
    for subject in target_train_subjects_list:
        for npz_name in os.listdir(os.path.join(target_path, subject)):
            abs_npz_name = os.path.join(target_path, subject, npz_name)
            target_train_npz_name_list.append(abs_npz_name)
        np.savetxt(os.path.join(txt_save_path, 'target_train.txt'), np.array(target_train_npz_name_list), fmt ="%s")
    #target的test
    target_test_npz_name_list = []
    for subject in target_test_subjects_list:
        for npz_name in os.listdir(os.path.join(target_path, subject)):
            abs_npz_name = os.path.join(target_path, subject, npz_name)
            target_test_npz_name_list.append(abs_npz_name)
        np.savetxt(os.path.join(txt_save_path, 'target_test.txt'), np.array(target_test_npz_name_list), fmt ="%s")
    print("subjects number:source:{}  target_train:{}   target_test:{}"
          .format(len(source_subjects_list),len(target_train_subjects_list),len(target_test_subjects_list)))
    print("2dslice number:source:{}  target_train:{}   target_test:{}".
          format(len(source_npz_name_list), len(target_train_npz_name_list), len(target_test_npz_name_list)))
if __name__ == '__main__':
    #先把BTCV的CT作为源数据
    # target_path="/home/liukai/AllData/AbdominalOrgansForDomain/ForSIFA/SIFAgf/CHAOS_MR_T2_npy"
    # source_path="/home/liukai/AllData/AbdominalOrgansForDomain/ForSIFA/SIFAgf/BTCV_CT_npy"
    # getTrainTest(source_path=source_path,target_path=target_path,
    #              txt_save_path='./')

    target_path="/home/liukai/AllData/liverTumorForDomain/ATLAS2023AfterProcess"
    source_path = "/home/liukai/AllData/liverTumorForDomain/LITSAfterProcess"
    getTrainTest(source_path=source_path,target_path=target_path,
                 txt_save_path='./LiverAndTumor')